
module LanguageTransform where
import Synchronizers
import LanguageDef
import Data.Set(Set)
import Data.Map(Map)
import qualified Data.Set as Set
import qualified Data.Map as Map
import Data.Maybe
import Data.List
import Debug.Trace
import Text.Printf


transform :: Node -> Either ShowS Node
transform node = --trace (show node ++ "\n\n") $  
    case (varCheck node >>> relationRefCheck node >>> mainCheck node) of
           Just msg -> Left msg
           Nothing -> --trace (show ((introduceLet) node)) $ 
                Right $  (introduceHaskellVal.introduceELet.introduceLet) node

mainCheck :: Node -> Maybe ShowS
mainCheck node = if Map.member "main" (getSigs node) 
    then Nothing
    else Just $ showString "There is no \"main\" method."

varCheck :: Node -> Maybe ShowS
varCheck node = varCheck' Set.empty node --maybeShow2String $ varCheck' Set.empty node

maybeShow2String :: Maybe ShowS -> Maybe String          
maybeShow2String msg = 
    case msg of
          Just showMsg -> Just $ showMsg ""
          Nothing -> Nothing



varCheck' :: Set String -> Node -> Maybe ShowS
varCheck' vars node = 
    (\(vars, maybeMsg) -> 
        foldl (\msg n -> msg >>> (varCheck' vars n)) 
            maybeMsg (getNodeChildren node))
        (varCheckSingle vars node)
        
varCheckSingle :: Set String -> Node -> (Set String, Maybe ShowS)
varCheckSingle vars Node{nodeType=VarNode varName, pos=pos} 
    | Set.member varName vars = (vars, Nothing)
    | otherwise = (vars, Just $ pos2String pos.showString ":".
        shows varName.showString " is not declared.") 
varCheckSingle vars node@Node{nodeType=RelationDefNode _}
    =  (Set.union ((varList2Set.head.getNodeChildren) node) vars, Nothing)
    where
        varList2Set node = Set.fromList (map var2string (getNodeChildren node))
varCheckSingle vars node@Node{nodeType=ForAllNode}
    = (Set.insert varName vars, Nothing)
    where
        varName = (var2string.head.tail.getNodeChildren) node
varCheckSingle vars node@Node{nodeType=ExistsNode}
    = (Set.insert varName vars, Nothing)
    where
        varName = (var2string.head.tail.getNodeChildren) node
varCheckSingle vars node@Node{nodeType=Exists'Node}
    = (Set.insert varName vars, Nothing)
    where
        varName = (var2string.head.tail.getNodeChildren) node
varCheckSingle vars node@Node{nodeType=OneNode}
    = (Set.insert varName vars, Nothing)
    where
        varName = (var2string.head.tail.getNodeChildren) node
varCheckSingle vars node@Node{nodeType=One'Node}
    = (Set.insert varName vars, Nothing)
    where
        varName = (var2string.head.tail.tail.getNodeChildren) node
varCheckSingle vars node@Node{nodeType=LetNode}
    = (Set.insert varName vars, Nothing)
    where
        varName = (var2string.head.getNodeChildren) node
varCheckSingle vars _ = (vars, Nothing)

var2string node@Node{nodeType=VarNode name} =  name


(>>>) :: Maybe ShowS -> Maybe ShowS -> Maybe ShowS
Just msg1 >>> Just msg2 = Just $ msg1.showString "\n".msg2
Just msg >>> Nothing = Just msg
Nothing >>> Just msg = Just msg
Nothing >>> Nothing = Nothing

checkConformance :: Node -> Map String Int -> Maybe ShowS
checkConformance node@Node{nodeType=RelationRefNode name, pos=pos} sigs 
    = case Map.lookup name sigs of
            Just paraNumber -> 
                let argNumber = (length.getNodeChildren) node in
                if paraNumber == argNumber 
                    then Nothing
                    else Just $ pos2String pos.
                        showString (printf "'%s' requires %d parameters, but is provided %d" name paraNumber argNumber) 
            Nothing -> Just $ pos2String pos.showString ":".shows name.
                            showString " is not declared."
checkConformance _ _ = Nothing

getSigs :: Node->Map String Int
getSigs node = foldl appendSingle Map.empty (getNodeChildren node)
    where 
        appendSingle sig Node{nodeType=RelationDefNode name, children=Node{children=vars}:_} =
            Map.insert name (length vars) sig
        appendSingle sig _ = sig

relationRefCheck :: Node -> Maybe ShowS
relationRefCheck node = relationRefCheck' (getSigs node) node -- maybeShow2String $ relationRefCheck' (getSigs node) node
relationRefCheck' sigs node = 
    foldl (\maybeMsg node->maybeMsg >>> relationRefCheck' sigs node) 
        (checkConformance node sigs) (getNodeChildren node)

isExprNode :: Node -> Bool
isExprNode node = 
    case getNodeType node of
          DictGetNode -> True 
          IsConstNode -> True
          ETraceNode _ -> True
          LetNode -> isExprNode (getNodeChildren node !! 2)
          _ -> False
          
emptyPos = Pos 0 0

introduceLet :: Node -> Node
introduceLet node = snd $ introduceLet' 0 node 

varNode :: Int -> Node
varNode varid = Node (VarNode ("@var" ++ show varid)) emptyPos []



introduceLet' :: Int -> Node -> (Int, Node)
introduceLet' varid node@Node{nodeType=LetNode, pos=pos, children=children} = passResult varid node
introduceLet' varid node@Node{nodeType=ETraceNode _, pos=pos, children=children} = passResult varid node
introduceLet' varid node@Node{nodeType=nodeType, pos=pos, children=children}
    | isJust $ find isExprNode children = 
        let Just i = findIndex isExprNode children
            varid' :: Int
            (varid', left) = ((foldl merge (varid + 1, [])).(take i)) children
            varid'' :: Int
            (varid'', right) = ((foldl merge (varid', [])).(drop $ i + 1)) children
            newchildren = left ++ [varNode varid] ++ right
            (rvarid, newexpr) = introduceLet' varid'' (children!!i)
            (rvarid', innerNode) = introduceLet' rvarid (Node nodeType pos newchildren)
        in
        (rvarid', Node LetNode emptyPos [varNode varid, newexpr, innerNode])
    | otherwise = passResult varid node
merge (varid, nodelist) node = (fst r, nodelist ++ [snd r]) where { r = introduceLet' varid node }
passResult varid node@Node{nodeType=nodeType, pos=pos, children=children} = let (rvarid, r) = foldl merge (varid, []) children in (rvarid, Node nodeType pos r) 
        
     


introduceELet :: Node -> Node
introduceELet node@Node{nodeType=nodeType, pos=pos, children=children}  
    | nodeType == LetNode && isExprNode (children !! 2) = Node LeteNode pos (map introduceELet children)  
    | otherwise = Node nodeType pos (map introduceELet children) 
    

introduceHaskellVal :: Node -> Node
introduceHaskellVal = introduceHaskellVal' Set.empty 
introduceHaskellVal' :: Set String -> Node -> Node
introduceHaskellVal' vars node@Node{nodeType=VarNode varName, pos=pos, children=children} 
    | Set.member varName vars = Node{nodeType=HaskellVarNode varName, pos=pos, children=children}
    | otherwise =  node
introduceHaskellVal' vars node@Node{nodeType=RelationDefNode name, pos=pos, children=[paraList, cstraint]}
    = Node (RelationDefNode name) pos [introduceHaskellVal' varNames paraList,introduceHaskellVal' varNames cstraint] 
    where
     varNames = Set.union (getVarNames paraList) vars
     getVarNames node = Set.fromList $ map var2string (getNodeChildren node)
introduceHaskellVal' vars node@Node{nodeType=ForAllNode, pos=pos, children=[dexpr, vvar, cstraint]}
    = Node ForAllNode pos [introduceHaskellVal' vars dexpr, vvar, introduceHaskellVal' (Set.delete (var2string vvar) vars) cstraint] 
introduceHaskellVal' vars node@Node{nodeType=ExistsNode, pos=pos, children=[dexpr, vvar, cstraint]}
    = Node ExistsNode pos [introduceHaskellVal' vars dexpr, vvar, introduceHaskellVal' (Set.delete (var2string vvar) vars) cstraint]
introduceHaskellVal' vars node@Node{nodeType=Exists'Node, pos=pos, children=[dexpr, vvar, cstraint]}
    = Node Exists'Node pos [introduceHaskellVal' vars dexpr, vvar, introduceHaskellVal' (Set.delete (var2string vvar) vars) cstraint]
introduceHaskellVal' vars node@Node{nodeType=OneNode, pos=pos, children=[dexpr, vvar, cstraint]}
    = Node OneNode pos [introduceHaskellVal' vars dexpr, vvar, introduceHaskellVal' (Set.delete (var2string vvar) vars) cstraint]
introduceHaskellVal' vars node@Node{nodeType=One'Node, pos=pos, children=[kexpr, dexpr, vvar, cstraint]}
    = Node One'Node pos [introduceHaskellVal' vars kexpr, introduceHaskellVal' vars dexpr, vvar, introduceHaskellVal' (Set.delete (var2string vvar) vars) cstraint]
introduceHaskellVal' vars node@Node{nodeType=LetNode, pos=pos, children=[vvar, expr, cstraint]}
    = Node LetNode pos [vvar, introduceHaskellVal' vars expr, introduceHaskellVal' (Set.delete (var2string vvar) vars) cstraint]
introduceHaskellVal' vars Node{nodeType=nodeType, pos=pos, children=children} 
    = Node nodeType pos (map (introduceHaskellVal' vars) children)  
